import logging
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
from tqdm import tqdm
import datetime
from pinecone import Pinecone, ServerlessSpec

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

digits = load_digits(n_class=10)
X = digits.data
y = digits.target

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

vectors_train = [(str(i), X_train[i].tolist(), {"label": int(y_train[i])}) for i in range(len(X_train))]

pinecone = Pinecone(api_key="816f2326-8a49-460e-8818-d810d393d713")

# 索引名称
index_name = "mnist-index"

existing_indexes = pinecone.list_indexes()

if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pinecone.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pinecone.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pinecone.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")

batch_size = 1000
for i in tqdm(range(0, len(vectors_train), batch_size), desc="Uploading training data"):
    batch = vectors_train[i:i + batch_size]
    index.upsert(batch)  # 实际应用中取消注释

logging.info(f"成功创建索引，并上传了{len(vectors_train)}条数据")

# 测试准确率
from collections import Counter
k = 11
correct_predictions = 0
for i in tqdm(range(len(X_test)), desc="Testing accuracy with k=11"):
    query_data = X_test[i].tolist()
    #results = index.query(vector=query_data, top_k=k, include_metadata=True)  # 实际应用中取消注释
    #mock_results for simulation 
    mock_results = [{'metadata': {'label': int(y_train[np.random.randint(0, len(y_train))])}} for _ in range(k)]
    labels = [match['metadata']['label'] for match in mock_results]
    prediction = Counter(labels).most_common(1)[0][0]
    if prediction == y_test[i]:
        correct_predictions += 1

accuracy = correct_predictions / len(X_test)
logging.info(f"当k={k}时，使用Pinecone的准确率为{accuracy:.4f}")